# -*- coding: utf-8 -*-
from sklearn.model_selection import train_test_split
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import os

model_file = 'my_FC.pth'

def evaluate_model(model, X_test, y_test, epoch, device):
    model.eval()  # 设置模型为评估模式
    with torch.no_grad():  # 不计算梯度
        inputs = torch.tensor(X_test, dtype=torch.float32).to(device)  # 移动到GPU
        outputs = model(inputs)  # 进行预测
        _, predicted = torch.max(outputs.data, 1)  # 获取预测的类别
        accuracy = (predicted.cpu().numpy() == y_test).mean()  # 将预测移动回CPU并计算准确率
        print('epoch:{} Accuracy: {}'.format(epoch, accuracy))  # 打印准确率


num_epochs = 500  # 或其他适当的数字

data = np.load('../data/data_mean_vec.npz')
sentences_array = data['sentences']
tags_array = data['tags']

X_train, X_test, y_train, y_test = train_test_split(sentences_array, tags_array, test_size=0.2, random_state=42)


class SimpleNN(nn.Module):
    def __init__(self, input_dim):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(input_dim, 128)  # 第一层，输入300维，输出128维
        self.fc2 = nn.Linear(128, 64)  # 第二层，输出64维
        self.fc3 = nn.Linear(64, 2)  # 输出层，2个类别（0或1）

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


model = SimpleNN(input_dim=300)
criterion = nn.CrossEntropyLoss()  # 使用交叉熵损失
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
if os.path.exists(model_file):
    model.load_state_dict(torch.load(model_file))
    print("模型权重已加载。")
else:
    print("模型文件不存在，无法加载。")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)  # 将模型移动到GPU

# 训练循环
for epoch in range(num_epochs):
    model.train()
    optimizer.zero_grad()

    # 转为Tensor并移动到GPU
    inputs = torch.tensor(X_train, dtype=torch.float32).to(device)
    labels = torch.tensor(y_train, dtype=torch.long).to(device)

    outputs = model(inputs)  # 前向传播
    loss = criterion(outputs, labels)  # 计算损失
    loss.backward()  # 反向传播
    optimizer.step()  # 更新参数

    # 在GPU上评估
    evaluate_model(model, X_test, y_test, epoch, device)
torch.save(model.state_dict(), model_file )
print("模型已保存为 {}".format(model_file))
